import pandas as pd
import torch
import numpy as np

# 确保在CPU上运行
device = torch.device("cpu")

def CROSS_TORCH(S1, S2):
    # 处理标量情况
    if np.isscalar(S2):
        S2 = np.full_like(S1, S2) if hasattr(S1, '__len__') else S2
    
    # 将数据移动到CPU
    S1_t = torch.tensor(S1.values if hasattr(S1, 'values') else S1, device=device)
    S2_t = torch.tensor(S2.values if hasattr(S2, 'values') else S2, device=device)
    
    # 确保张量维度正确
    if S1_t.dim() == 0:
        S1_t = S1_t.unsqueeze(0)
    if S2_t.dim() == 0:
        S2_t = S2_t.unsqueeze(0)
    
    # 检查数据长度是否足够
    if len(S1_t) < 2 or len(S2_t) < 2:
        # 返回全False的结果，长度与输入一致
        result = torch.zeros_like(S1_t, dtype=torch.bool)
        return pd.Series(result.cpu().numpy(), index=S1.index if hasattr(S1, 'index') else None)

    # 计算金叉
    shifted_S1 = S1_t[:-1]
    shifted_S2 = S2_t[:-1]
    result = torch.empty_like(S1_t, dtype=torch.bool)
    result[0] = False
    result[1:] = ~(shifted_S1 > shifted_S2) & (S1_t[1:] > S2_t[1:])
    
    return pd.Series(result.cpu().numpy(), index=S1.index if hasattr(S1, 'index') else None)

def EXPMEMA(s, n, m=4):
    s = pd.Series(s)
    alpha = 2 / (n + 1)
    # EMA 初始值处理：通常 EMA 的第一个值就是第一个数据点，或者取前N个数据的均值
    # 这里保持原有逻辑，用rolling(m).mean()，但要注意可能产生NaN，需要填充
    ema = s.ewm(span=n, adjust=False).mean() # 使用 span=n 更符合通达信的EMA平滑系数
    return ema.values

def CROSS_DOWN_TORCH(S1, S2):
    # 处理标量情况
    if np.isscalar(S2):
        S2 = np.full_like(S1, S2) if hasattr(S1, '__len__') else S2
    
    # 将数据移动到CPU
    S1_t = torch.tensor(S1.values if hasattr(S1, 'values') else S1, device=device)
    S2_t = torch.tensor(S2.values if hasattr(S2, 'values') else S2, device=device)
    
    # 确保张量维度正确
    if S1_t.dim() == 0:
        S1_t = S1_t.unsqueeze(0)
    if S2_t.dim() == 0:
        S2_t = S2_t.unsqueeze(0)

    # 检查数据长度是否足够
    if len(S1_t) < 2 or len(S2_t) < 2:
        result = torch.zeros_like(S1_t, dtype=torch.bool)
        return pd.Series(result.cpu().numpy(), index=S1.index if hasattr(S1, 'index') else None)
    
    # 计算下穿
    shifted_S1 = S1_t[:-1]
    shifted_S2 = S2_t[:-1]
    result = torch.empty_like(S1_t, dtype=torch.bool)
    result[0] = False
    result[1:] = (shifted_S1 > shifted_S2) & (S1_t[1:] < S2_t[1:])
    
    return pd.Series(result.cpu().numpy(), index=S1.index if hasattr(S1, 'index') else None)

# 注意：CROSS_TORCH_DOWN 和 CROSS_TORCH 逻辑重复，且命名容易混淆。
# 如果 CROSS_TORCH 已经用于金叉，那么 CROSS_DOWN_TORCH 用于死叉更清晰。
# 这里的 CROSS_TORCH_DOWN 与 CROSS_TORCH 逻辑完全一样，我将保留它，但建议检查你的实际需求。
def CROSS_TORCH_DOWN (S1, S2): # 建议重命名或删除，因为它与 CROSS_TORCH 相同
    # 处理标量情况
    if np.isscalar(S2):
        S2 = np.full_like(S1, S2) if hasattr(S1, '__len__') else S2
    
    # 将数据移动到CPU
    S1_t = torch.tensor(S1.values if hasattr(S1, 'values') else S1, device=device)
    S2_t = torch.tensor(S2.values if hasattr(S2, 'values') else S2, device=device)
    
    # 确保张量维度正确
    if S1_t.dim() == 0:
        S1_t = S1_t.unsqueeze(0)
    if S2_t.dim() == 0:
        S2_t = S2_t.unsqueeze(0)
    
    # 检查数据长度是否足够
    if len(S1_t) < 2 or len(S2_t) < 2:
        # 返回全False的结果，长度与输入一致
        result = torch.zeros_like(S1_t, dtype=torch.bool)
        return pd.Series(result.cpu().numpy(), index=S1.index if hasattr(S1, 'index') else None)
    
    # 计算金叉
    shifted_S1 = S1_t[:-1]
    shifted_S2 = S2_t[:-1]
    result = torch.empty_like(S1_t, dtype=torch.bool)
    result[0] = False
    result[1:] = ~(shifted_S1 > shifted_S2) & (S1_t[1:] > S2_t[1:])
    
    return pd.Series(result.cpu().numpy(), index=S1.index if hasattr(S1, 'index') else None)


def FILTER_TORCH(S, N):
    # 将输入转换为 PyTorch 张量（如果还不是）
    if not isinstance(S, torch.Tensor):
        S = torch.tensor(S, dtype=torch.bool, device=device) # 确保在CPU上
    
    # 创建一个与 S 相同大小的零张量
    result = torch.zeros_like(S, dtype=torch.bool, device=device) # 确保在CPU上
    
    # 找到所有满足条件的索引
    trigger_indices = torch.where(S)[0]
    
    # 对于每个触发点，设置其后 N 个位置为 True
    for idx in trigger_indices:
        start = idx + 1
        end = min(start + N, len(S))
        result[start:end] = True
    
    # 原始条件与"未被过滤"的位置做与操作
    return S & (~result)

# 计算动态移动平均 MM， 优化向量化实现（减少显存）
def DMA_torch(S, A):
    # 确保在CPU上运行
    S_tensor = torch.as_tensor(S.values if isinstance(S, pd.Series) else S, dtype=torch.float32, device=device)
    A_tensor = torch.as_tensor(A.values if isinstance(A, pd.Series) else A, dtype=torch.float32, device=device)
    
    # 使用累积乘积替代显存密集型矩阵
    cumprod = torch.cumprod(1 - A_tensor, dim=0)
    weighted_S = S_tensor * A_tensor
    result = torch.cumsum(weighted_S * cumprod, dim=0) / cumprod
    return result.cpu().numpy()

def XMA(S, N):
    '''
    正确版 XMA 超平滑均线
    每日基于前一日XMA和当日价格修正，防止累加放大。
    '''
    S = np.array(S)
    alpha = 2 / (N + 1)

    XMA_arr = np.zeros_like(S, dtype=float) # 明确数据类型为float
    if len(S) > 0:
        XMA_arr[0] = S[0]  # 初始取第一天价格

        for i in range(1, len(S)):
            XMA_arr[i] = XMA_arr[i-1] + alpha * (S[i] - XMA_arr[i-1])

    return XMA_arr